# -*- coding: utf-8 -*-
"""
Created on Fri Sep 27 17:08:30 2019

@author: swls
"""

import os
import sys
import cv2
import json
import imageio
import numpy as np
import imgaug as ia
import imgaug.augmenters as iaa
from imgaug.augmentables.polys import Polygon, PolygonsOnImage

def read_poloygon_json_to_list(data_dir):
    with open(data_dir + os.sep + "via_region_data.json","r") as rf:
        data_dict = json.loads(rf.read())
    data_new =  {}
    for filename in data_dict.keys():
        xys = []
        for one_label in data_dict[filename]["regions"]:
            xs = one_label["shape_attributes"]["all_points_x"]
            ys = one_label["shape_attributes"]["all_points_y"]
            xy = []
            for x,y in zip(xs,ys):
                xy.append((x,y))
            xys.append(xy)
        data_new[data_dir + os.sep + data_dict[filename]["filename"]] = xys
    #print(len(data_new.keys()))
    return data_new

def read_rect_json_to_list(data_dir):
    with open(data_dir + os.sep + "via_region_data.json","r") as rf:
        data_dict = json.loads(rf.read())
    data_new =  {}
    for filename in data_dict.keys():
        xys = []
        for one_label in data_dict[filename]["regions"]:
            x0 = one_label["shape_attributes"]["x"]
            y0 = one_label["shape_attributes"]["y"]
            x1 = one_label["shape_attributes"]["width"] + x0
            y1 = one_label["shape_attributes"]["height"] + y0
            xy = [x0,y0,x1,y1]
            xys.append(xy)
        data_new[data_dir + os.sep + data_dict[filename]["filename"]] = xys
    #print(len(data_new.keys()))
    return data_new

def convert_to_json_and_save_img(save_dir, image_aug, polygon_aug):

    via_dict = dict()
    for idx, (image, polygons) in enumerate(zip(image_aug, polygon_aug)):
        name = "imgaug_" + str(idx) + ".jpg"
        image_path = save_dir + os.sep + name
        imageio.imwrite(image_path, image)
        current_dict = dict()
        current_dict["filename"] = name
        regions = []

        for plg in polygons:
            shape_attributes = dict()
            shape_attributes["name"] = 'polygon'
            shape_attributes["all_points_x"] = [int(x) for x in plg.xx]
            shape_attributes["all_points_y"] = [int(y) for y in plg.yy]
            regions.append({"shape_attributes": shape_attributes, "region_attributes":{"type":"text"}})
        current_dict["regions"] = regions
        current_dict["file_attributes"] = {}
        via_dict[name] = current_dict

    with open(save_dir + os.sep + "via_region_data_ori.json", "w") as rf:
        rf.write(json.dumps(via_dict))

def convert_to_rect_json_and_save_img(save_dir, image_aug, polygon_aug):

    via_dict = dict()
    for idx, (image, polygons) in enumerate(zip(image_aug, polygon_aug)):
        name = "imgaug_" + str(idx) + ".jpg"
        image_path = save_dir + os.sep + name
        imageio.imwrite(image_path, image)
        current_dict = dict()
        current_dict["filename"] = name
        regions = []

        for plg in polygons:
            shape_attributes = dict()
            shape_attributes["name"] = 'rect'
            shape_attributes["x"] = int(np.min([plg.x1_int, plg.x2_int]))
            shape_attributes["y"] = int(np.min([plg.y1_int, plg.y2_int]))
            shape_attributes["width"] = int(abs(plg.x1_int- plg.x2_int))
            shape_attributes["height"] = int(abs(plg.y1_int- plg.y2_int))

            regions.append({"shape_attributes": shape_attributes, "region_attributes":{"type":"rect"}})
        current_dict["regions"] = regions
        current_dict["file_attributes"] = {}
        via_dict[name] = current_dict

    with open(save_dir + os.sep + "via_region_data_ori.json", "w") as rf:
        rf.write(json.dumps(via_dict))

def aug_image(data_dict):
    io_images = []
    polygons = []

    for idx, image_path in enumerate(data_dict.keys()):
        img = imageio.imread(image_path)
        new_img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
        io_images.append(new_img)

        # polygons.append([ ia.Polygon([(x,y) for x,y in xy] )  for xy in data_dict[image_path]])
        polygons.append([ ia.BoundingBox(x1=boxes[0], y1=boxes[1],x2=boxes[2],y2=boxes[3]) for boxes in  data_dict[image_path] ])
        # if idx == 2:
        #     break
    # new_image = ia.imresize_single_image(io_image, 0.5)

    sometimes = lambda aug: iaa.Sometimes(0.5, aug)

    seq = iaa.Sequential(
        [
            iaa.Rot90(1, keep_size=False)
            # iaa.AdditiveGaussianNoise(scale=(0.01 * 255)),
            # iaa.Affine(translate_px={"x":(1,5)}),
            # sometimes([

            # ]),
            # iaa.OneOf(
            #     [
            #         iaa.Dropout(p=(0.01, 0.1), per_channel=0.5),
            #         #iaa.Flipud(0.5),
            #         # iaa.LinearContrast((0.5,1.5), per_channel=0.5),
            #         iaa.Add((-10, 10), per_channel=0.5),
            #         #iaa.SimplexNoiseAlpha(iaa.OneOf([
            #         #    iaa.EdgeDetect(alpha=(0.5, 1.0)),
            #         #    iaa.DirectedEdgeDetect(alpha=(0.5, 1.0), direction=(0.0, 1.0))
            #         #])),

            #     ]
            # ),
            # iaa.OneOf([
            #     iaa.GaussianBlur((0, 3.0)),
            #     iaa.MedianBlur(k=(3,5)),
            #     iaa.AverageBlur(k=(3,5))
            # ])
        ],
        random_order=True
    )

    # image_aug, polygon_aug = seq(images=io_images, polygons=polygons)
    image_aug, polygon_aug = seq(images=io_images, bounding_boxes=polygons)
    return image_aug, polygon_aug

    # ia.imshow(new_image)
    

def main(data_dir):

    data_dict = read_rect_json_to_list(data_dir)

    ori_dict = {}
    aug_dict = {}
    for path in data_dict.keys():
        value = os.path.basename(path).split("_")[0]
        if 5 <= int(value) <= 20:
            aug_dict[path] = data_dict[path]
        else:
            ori_dict[path] = data_dict[path]

    image_aug, polygon_aug = aug_image(aug_dict)

    new_save_dir = data_dir + "_imgaug"
    if not os.path.exists(new_save_dir):
        os.mkdir(new_save_dir)
    # convert_to_json_and_save_img(new_save_dir, image_aug, polygon_aug)

    for path in ori_dict.keys():
        img = imageio.imread(path)
        poly = [ ia.BoundingBox(x1=b[0], y1=b[1],x2=b[2], y2=b[3]) for b in  ori_dict[path] ]
        image_aug.append(img)
        polygon_aug.append(poly)
    convert_to_rect_json_and_save_img(new_save_dir, image_aug, polygon_aug)

if __name__ == "__main__":
    # data_dir = "C:\\Users\\swls\\Desktop\\being\\2019-09-24\\label"

    #5 - 20
    data_dir = r"C:\Users\swls\Documents\WeChat Files\A360697611\FileStorage\File\2019-11\zj_image"
    main(data_dir)


